import logging
import torch
import torch.nn as nn
from mmcv.cnn import VGG, constant_init, kaiming_init, normal_init, xavier_init
from mmcv.runner import load_checkpoint
from ..module import BACKBONE


@BACKBONE.register_module
class Vgg(VGG):
    """VGG Backbone network for single-shot-detection

    Args:
        input_size (int): width and height of input, from {300, 512}.
        depth (int): Depth of vgg, from {11, 13, 16, 19}.
        out_indices (Sequence[int]): Output from which stages.
    """

    def __init__(self,
                 depth,
                 with_last_pool=False,
                 ceil_mode=True,
                 out_indices=(3,),
                 pretrained=None
                 ):
        # TODO: in_channels for mmcv.VGG
        super(Vgg, self).__init__(
            depth,
            with_last_pool=with_last_pool,
            ceil_mode=ceil_mode,
            out_indices=out_indices)
        self.pretrained = pretrained

    def init_weights(self):
        pretrained = self.pretrained
        if isinstance(pretrained, str):
            logger = logging.getLogger()
            load_checkpoint(self, pretrained, strict=False, logger=logger)
        elif pretrained is None:
            for m in self.features.modules():
                if isinstance(m, nn.Conv2d):
                    kaiming_init(m)
                elif isinstance(m, nn.BatchNorm2d):
                    constant_init(m, 1)
                elif isinstance(m, nn.Linear):
                    normal_init(m, std=0.01)
        else:
            raise TypeError('pretrained must be a str or None')

    def forward(self, x):
        outs = []
        for i, layer in enumerate(self.features):
            x = layer(x)
        outs.append(x)
        if len(outs) > 1:
            return tuple(outs)
        return outs[0]